import argparse
import os
from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import mannwhitneyu

class CellViabilityAnalysis:
    """
    用于分析化合物对细胞系活性影响的类，使用曲线下面积（AUC）方法。

    该类提供工具来计算总AUC和分段AUC，处理剂量-响应曲线和时间序列数据，
    执行统计比较，计算IC50值，并生成出版质量的图表。

    该类可应用于时间依赖性实验（同一浓度不同时间点）和
    剂量依赖性实验（不同浓度同一时间点）的数据分析，支持
    细胞株特异性分析和多化合物比较。
    """

    def __init__(self, data, sample_id_col, group_col, x_col, response_col, analysis_type='time', cell_line_col=None):
        """使用细胞活性数据初始化分析对象。

        初始化时会自动查找与提供的列名匹配的实际列，支持大小写不敏感的匹配。

        Args:
            data (pd.DataFrame): 包含细胞活性数据的DataFrame。
            sample_id_col (str): 包含唯一样本标识符的列名（如孔位ID）。
            group_col (str): 指定实验组的列名（如化合物名称）。
            x_col (str): 包含自变量的列名（时间或浓度）。
            response_col (str): 包含细胞活性测量值的列名。
            analysis_type (str, optional): 分析类型，'time'（时间依赖性）或'dose'（剂量依赖性）。
                默认值为'time'。
            cell_line_col (str, optional): 包含细胞株信息的列名。如果为None，则尝试自动
                查找包含'cell'的列。
        """
        # 创建列名的小写映射，用于大小写不敏感的比较
        data_columns_lower = {col.lower(): col for col in data.columns}
        
        # 查找列名的函数
        def find_column(keyword, required=True):
            for col_lower, col in data_columns_lower.items():
                if keyword.lower() in col_lower:
                    return col
            if required:
                raise ValueError(f"找不到包含'{keyword}'的列")
            return None
        
        # 查找各列
        actual_sample_id_col = find_column(sample_id_col)
        actual_group_col = find_column(group_col)
        
        # 特殊处理x_col和response_col：如果分析类型是dose且关键词包含concentration/viability，
        # 则进行更宽松的匹配
        if analysis_type == 'dose' and 'concentration' in x_col.lower():
            actual_x_col = find_column('concentration')
        else:
            actual_x_col = find_column(x_col)
            
        if 'viability' in response_col.lower():
            actual_response_col = find_column('viability')
        else:
            actual_response_col = find_column(response_col)
        
        # 查找细胞株列（如果提供）
        actual_cell_line_col = None
        if cell_line_col:
            try:
                actual_cell_line_col = find_column(cell_line_col)
            except ValueError:
                # 如果找不到细胞株列，尝试查找包含'cell'的列
                try:
                    actual_cell_line_col = find_column('cell', required=False)
                except:
                    pass

        self.data = data.copy()
        self.sample_id_col = actual_sample_id_col
        self.group_col = actual_group_col
        self.x_col = actual_x_col
        self.response_col = actual_response_col
        self.analysis_type = analysis_type
        self.cell_line_col = actual_cell_line_col
        
        # 获取所有唯一的组名
        self.groups = self.data[self.group_col].unique()
        
        # 设置默认调色板并创建组名到颜色的映射
        colors = sns.color_palette("husl", len(self.groups))
        self.color_map = {group: color for group, color in zip(self.groups, colors)}
        
        # 保留palette属性以保持向后兼容
        self.palette = colors
        
        self.auc_results = None

    def _calculate_single_auc(self, sample_df, total_x_range=None):
        """计算单个样本的曲线下面积(AUC)。

        使用梯形法则计算指定样本数据的曲线下面积，并能处理不完整数据。
        对于剂量依赖性分析，会特别处理确保包含0浓度点和最大浓度点。

        Args:
            sample_df (pd.DataFrame): 包含单个样本数据的DataFrame。
            total_x_range (tuple, optional): 自变量的总范围，格式为(min_value, max_value)。
                默认为None。

        Returns:
            float: 计算得到的AUC值。
        """
        sample_df = sample_df.sort_values(by=self.x_col).reset_index()
        
        x_vals = sample_df[self.x_col].values
        response_vals = sample_df[self.response_col].values
        
        # 对于剂量依赖性分析，可能需要特殊处理
        if self.analysis_type == 'dose' and total_x_range is not None:
            # 确保包含0浓度点
            min_x = min(total_x_range)
            if min_x not in x_vals:
                x_vals = np.append(min_x, x_vals)
                # 假设0浓度时细胞活性为100%（可以根据实际情况调整）
                response_vals = np.append(100.0, response_vals)
            
            # 确保包含最大浓度点
            max_x = max(total_x_range)
            if max_x not in x_vals:
                # 对最大浓度点进行延伸或插值
                x_vals = np.append(x_vals, max_x)
                # 使用最后已知值延伸
                response_vals = np.append(response_vals, response_vals[-1])
        
        return np.trapezoid(response_vals, x_vals)

    def calculate_auc(self, total_x_range=None, by_cell_line=False):
        """计算指定范围内每个样本的总AUC。

        根据分析类型执行不同的计算策略：
        - 对于时间依赖性分析，按组分组，每个组作为一个样本计算AUC。
        - 对于剂量依赖性分析，可选择按样本ID或按细胞株和化合物分组计算。

        Args:
            total_x_range (tuple, optional): 自变量的总范围，格式为(min_value, max_value)。
                用于纠正不完整数据。对于剂量依赖性分析尤为重要。
                如果未提供，将使用数据中的实际最小和最大值。
            by_cell_line (bool, optional): 是否按细胞株分组计算AUC。
                当data中包含cell_line_col时有效。默认为False。

        Returns:
            pd.DataFrame: 包含计算结果的DataFrame，包括：
                - sample_id: 样本ID（剂量依赖性分析）
                - group: 实验组名称
                - cell_line: 细胞株名称（如果适用）
                - total_auc: 计算得到的总AUC值
        """
        auc_data = []
        
        # 确定总体范围（如果未提供）
        if total_x_range is None:
            total_x_range = (self.data[self.x_col].min(), self.data[self.x_col].max())
        
        if self.analysis_type == 'time':
            # 对于时间分析，按group分组，每个组作为一个样本计算AUC
            for group, group_df in self.data.groupby(self.group_col):
                # 对每个时间点计算平均值，创建平均响应曲线
                avg_group_df = group_df.groupby(self.x_col)[self.response_col].mean().reset_index()
                auc = self._calculate_single_auc(avg_group_df, total_x_range)
                auc_data.append({
                    self.group_col: group,
                    'total_auc': auc
                })
        else:
            # 对于剂量分析
            if by_cell_line and self.cell_line_col:
                # 按细胞株和组（化合物）分组计算AUC
                for (cell_line, group), subset_df in self.data.groupby([self.cell_line_col, self.group_col]):
                    # 确保数据按x值（浓度）排序
                    subset_df = subset_df.sort_values(by=self.x_col)
                    auc = self._calculate_single_auc(subset_df, total_x_range)
                    auc_data.append({
                        self.cell_line_col: cell_line,
                        self.group_col: group,
                        'total_auc': auc
                    })
            else:
                # 原始行为：按sample_id分组
                for sample_id, group_df in self.data.groupby(self.sample_id_col):
                    auc = self._calculate_single_auc(group_df, total_x_range)
                    group = group_df[self.group_col].iloc[0]
                    
                    # 如果数据中包含细胞株信息，也提取出来
                    cell_line_info = None
                    if self.cell_line_col in group_df.columns:
                        cell_line_info = group_df[self.cell_line_col].iloc[0]
                    
                    row_data = {
                        self.sample_id_col: sample_id,
                        self.group_col: group,
                        'total_auc': auc
                    }
                    
                    if cell_line_info is not None:
                        row_data[self.cell_line_col] = cell_line_info
                    
                    auc_data.append(row_data)
        
        self.auc_results = pd.DataFrame(auc_data)
        return self.auc_results

    def compare_groups(self, metric, group1, group2, test='mann-whitneyu'):
        """执行统计测试以比较两组之间的指标。

        从已计算的结果中提取两组数据，并执行指定的统计测试进行比较。
        支持的指标包括'total_auc'和分段AUC值。

        Args:
            metric (str): 要比较的指标名称（例如，'total_auc'）。
            group1 (str): 第一组的名称。
            group2 (str): 第二组的名称。
            test (str, optional): 要使用的统计测试。目前仅支持'mann-whitneyu'。
                默认为'mann-whitneyu'。

        Returns:
            dict: 包含测试结果的字典，包括：
                - p_value: 计算得到的p值
                - statistic: 统计测试的统计量值
                - note: 如果数据不足等情况的备注信息（可选）

        Raises:
            ValueError: 如果指定的指标未计算或未找到。
            NotImplementedError: 如果指定的统计测试未实现。
        """
        if metric == 'total_auc' and self.auc_results is not None:
            data_source = self.auc_results
        elif metric.startswith('auc_') and self.segmented_auc_results is not None:
            data_source = self.segmented_auc_results
        else:
            raise ValueError(
                f"未找到指标 '{metric}'。请先运行相应的计算方法。"
            )

        vals1 = data_source[data_source[self.group_col] == group1][metric].dropna()
        vals2 = data_source[data_source[self.group_col] == group2][metric].dropna()

        if len(vals1) < 1 or len(vals2) < 1:
            return {'p_value': np.nan, 'statistic': np.nan, 'note': '比较数据不足'}

        if test == 'mann-whitneyu':
            stat, p_val = mannwhitneyu(vals1, vals2, alternative='two-sided')
            return {'p_value': p_val, 'statistic': stat}
        else:
                raise NotImplementedError(
                    f"测试 '{test}' 尚未实现。"
                )

    def plot_response_curves(self, ax=None, save_path=None, log_scale=False, **kwargs):
        """绘制每组的平均响应曲线。

        使用seaborn绘制带标准误差的平均响应曲线，支持时间依赖性和剂量依赖性数据。
        对于剂量依赖性数据，可选择使用对数刻度。

        Args:
            ax (matplotlib.axes.Axes, optional): 用于绘图的现有轴。如果为None，
                将创建新的图形和轴。默认为None。
            save_path (str, optional): 保存图形的路径。如果提供，将保存图形到指定路径。
                默认为None。
            log_scale (bool, optional): 是否对x轴使用对数刻度。剂量依赖性分析时通常使用。
                默认为False。
            **kwargs: 传递给seaborn.lineplot的其他关键字参数。

        Returns:
            matplotlib.axes.Axes: 包含绘制图形的轴对象。
        """
        if ax is None:
            fig, ax = plt.subplots(figsize=(8, 6))
        
        sns.lineplot(
            data=self.data,
            x=self.x_col,
            y=self.response_col,
            hue=self.group_col,
            errorbar='se',  # 标准误差
            ax=ax,
            palette=self.color_map,
            **kwargs
        )
        
        # 设置轴标签和标题
        if self.analysis_type == 'time':
            ax.set_title('Time-Dependent Cell Viability')
            ax.set_xlabel('Time (hours)')
        else:  # dose
            ax.set_title('Dose-Response Curve')
            ax.set_xlabel(f'Concentration')
            # 如果需要对数刻度
            if log_scale:
                ax.set_xscale('log')
                ax.set_xlabel(f'Concentration (log scale)')
        
        ax.set_ylabel(f'Cell Viability ({self.response_col})')
        ax.legend(title='Group')
        
        if save_path:
            plt.savefig(save_path, bbox_inches='tight', dpi=300)
            
        return ax

    def plot_auc_summary(self, metric, control_group=None, plot_type='box', ax=None, 
                        save_path=None, by_cell_line=False, selected_cell_lines=None, 
                        selected_compounds=None, **kwargs):
        """为所有组生成AUC指标的汇总图。

        支持多种图表类型，可按细胞株分组显示，并能添加统计显著性注释。
        对于大量数据，可通过筛选参数限制显示的细胞株或化合物。

        Args:
            metric (str): 要绘制的指标（'total_auc'或段AUC名称）。
            control_group (str, optional): 要比较的对照组。如果提供，
                将添加与对照组相比的统计显著性注释。默认为None。
            plot_type (str, optional): 图类型，支持：
                - 'box': 箱型图
                - 'bar': 柱状图
                - 'strip': 散点图
                - 'violin': 小提琴图
                - 'swarm': 蜂群图
                - 'heatmap': 热力图（仅在by_cell_line=True时有效）
                默认为'box'。
            ax (matplotlib.axes.Axes, optional): 用于绘图的现有轴。如果为None，
                将创建新的图形和轴。默认为None。
            save_path (str, optional): 保存图形的路径。如果提供，将保存图形到指定路径。
                默认为None。
            by_cell_line (bool, optional): 是否按细胞株展示结果。默认为False。
            selected_cell_lines (list, optional): 要显示的特定细胞株列表。
                如果为None，显示所有细胞株。默认为None。
            selected_compounds (list, optional): 要显示的特定化合物列表。
                如果为None，显示所有化合物。默认为None。
            **kwargs: 传递给seaborn绘图函数的其他关键字参数。

        Returns:
            matplotlib.axes.Axes or matplotlib.figure.Figure: 包含绘制图形的轴对象或图形对象。
                对于catplot类型的图表，返回整个图形。

        Raises:
            ValueError: 如果指定了不支持的图表类型。
        """
        # 确定数据源
        if metric == 'total_auc' and self.auc_results is not None:
            data_source = self.auc_results.copy()
        elif metric.startswith('auc_') and self.segmented_auc_results is not None:
            data_source = self.segmented_auc_results.copy()
        else:
            raise ValueError(f"未计算指标 '{metric}'。请先运行相应的计算方法。")

        # 筛选特定的细胞株或化合物
        if selected_cell_lines is not None and self.cell_line_col in data_source.columns:
            # 确保selected_cell_lines是列表而不是NumPy数组
            if hasattr(selected_cell_lines, 'tolist'):
                selected_cell_lines = selected_cell_lines.tolist()
            data_source = data_source[data_source[self.cell_line_col].isin(selected_cell_lines)]
        if selected_compounds is not None:
            data_source = data_source[data_source[self.group_col].isin(selected_compounds)]

        if ax is None:
            # 根据数据量动态调整图形大小
            if by_cell_line and self.cell_line_col in data_source.columns:
                cell_count = len(data_source[self.cell_line_col].unique())
                fig_width = max(8, cell_count * 0.5)
                fig, ax = plt.subplots(figsize=(fig_width, 6))
            else:
                group_count = len(data_source[self.group_col].unique())
                fig_width = max(8, group_count * 0.8)
                fig, ax = plt.subplots(figsize=(fig_width, 6))
        
        # 确定绘图方式
        if by_cell_line and self.cell_line_col in data_source.columns:
            # 按细胞株和化合物分组绘图
            if plot_type == 'heatmap':
                # 创建热力图数据
                pivot_data = data_source.pivot(index=self.cell_line_col, 
                                             columns=self.group_col, 
                                             values=metric)
                if 'cmap' not in kwargs:
                    kwargs['cmap'] = 'viridis'
                sns.heatmap(pivot_data, annot=False, fmt=".1f", ax=ax, **kwargs)
                ax.set_title(f'{metric} by Cell Line and Compound')
            else:
                # 分组条形图或其他类型
                if plot_type == 'bar':
                    sns.barplot(
                        data=data_source, 
                        x=self.cell_line_col, 
                        y=metric, 
                        hue=self.group_col,
                        ax=ax, 
                        palette=self.color_map,
                        **kwargs
                    )
                else:
                    sns.catplot(
                        data=data_source, 
                        x=self.cell_line_col, 
                        y=metric, 
                        hue=self.group_col,
                        kind=plot_type,
                        height=6,
                        aspect=max(2, len(data_source[self.cell_line_col].unique())/6),
                        palette=self.color_map,
                        **kwargs
                    )
                    plt.tight_layout()
                    return plt.gcf()  # 返回整个图形而不是轴
        else:
            # 原始行为：仅按组绘图
            if plot_type == 'box':
                sns.boxplot(
                    data=data_source, 
                    x=self.group_col, 
                    y=metric, 
                    hue=self.group_col, 
                    dodge=False, 
                    ax=ax, 
                    palette=self.color_map, 
                    legend=False,
                    **kwargs
                )
                sns.stripplot(data=data_source, x=self.group_col, y=metric, ax=ax, color='black', alpha=0.5)
            elif plot_type == 'bar':
                sns.barplot(
                    data=data_source, 
                    x=self.group_col, 
                    y=metric, 
                    hue=self.group_col,
                    dodge=False,
                    errorbar='sd', 
                    ax=ax, 
                    palette=self.color_map,
                    legend=False,
                    **kwargs
                )
                sns.stripplot(data=data_source, x=self.group_col, y=metric, ax=ax, color='black', alpha=0.5)
            elif plot_type == 'strip':
                sns.stripplot(
                    data=data_source, 
                    x=self.group_col, 
                    y=metric, 
                    hue=self.group_col,
                    ax=ax, 
                    palette=self.color_map,
                    legend=False,
                    **kwargs
                )
            elif plot_type == 'violin':
                sns.violinplot(
                    data=data_source, 
                    x=self.group_col, 
                    y=metric, 
                    hue=self.group_col,
                    dodge=False,
                    ax=ax, 
                    palette=self.color_map,
                    legend=False,
                    **kwargs
                )
                sns.stripplot(data=data_source, x=self.group_col, y=metric, ax=ax, color='black', alpha=0.5)
            elif plot_type == 'swarm':
                sns.swarmplot(
                    data=data_source, 
                    x=self.group_col, 
                    y=metric, 
                    hue=self.group_col,
                    ax=ax, 
                    palette=self.color_map,
                    legend=False,
                    **kwargs
                )
            else:
                raise ValueError(
                    f"不支持的图表类型: {plot_type}。"
                    f"支持的类型包括: 'box', 'bar', 'strip', 'violin', 'swarm'"
                )

        # 设置标题和标签
        if metric == 'total_auc':
            if self.analysis_type == 'time':
                base_title = 'Comparison of Total AUC (Time-Dependent)'
            else:
                base_title = 'Comparison of Total AUC (Dose-Response)'
        else:
            base_title = f'Comparison of {metric}'
            
        if by_cell_line and self.cell_line_col in data_source.columns:
            ax.set_title(f'{base_title} by Cell Line')
            ax.set_xlabel('Cell Line')
            # 旋转x轴标签以避免重叠
            plt.xticks(rotation=45, ha='right')
        else:
            ax.set_title(base_title)
            ax.set_xlabel('Group')
            
        ax.set_ylabel(f'{metric}')
        
        # 如果使用分组绘图，调整图例
        if by_cell_line and self.cell_line_col in data_source.columns and plot_type != 'heatmap':
            ax.legend(title='Compound', bbox_to_anchor=(1.05, 1), loc='upper left')
            plt.tight_layout()
            fig.subplots_adjust(right=0.85)  # 为图例留出空间

        # 如果指定了对照组，则添加统计注释
        if control_group is not None and not by_cell_line:
            groups = [g for g in data_source[self.group_col].unique() if g != control_group]
            y_max = data_source[metric].max()
            y_min = data_source[metric].min()
            y_range = y_max - y_min
            y_step = y_range * 0.1 if y_range > 0 else 1
            y_pos = y_max + y_step

            for i, group in enumerate(groups):
                res = self.compare_groups(metric, control_group, group)
                p_val = res['p_value']
                
                if p_val < 0.001:
                    significance_text = '***'
                elif p_val < 0.01:
                    significance_text = '**'
                elif p_val < 0.05:
                    significance_text = '*'
                else:
                    significance_text = 'ns'  # 无显著性差异（not significant）
                
                # 获取组的x坐标
                x1 = list(data_source[self.group_col].unique()).index(control_group)
                x2 = list(data_source[self.group_col].unique()).index(group)
                
                ax.plot([x1, x1, x2, x2], [y_pos, y_pos + y_step * 0.2, y_pos + y_step * 0.2, y_pos], 
                        lw=1.5, c='black')
                ax.text((x1 + x2) * 0.5, y_pos + y_step * 0.2, significance_text, ha='center', va='bottom', color='black')
                y_pos += y_step * 1.5

        if save_path:
            plt.savefig(save_path, bbox_inches='tight', dpi=300)
            
        return ax



def set_palette(analysis, palette_name, group_count):
    """设置分析对象的颜色映射。

    根据指定的调色板名称为分析对象创建组名到颜色的映射。
    支持使用seaborn内置调色板或自定义调色板。

    Args:
        analysis (CellViabilityAnalysis): 要设置调色板的分析对象。
        palette_name (str): 调色板名称或'custom'。
            可以是任何seaborn支持的调色板名称，如'husl', 'tab10', 'Set2'等。
            如果为'custom'，将使用预定义的自定义颜色列表。
        group_count (int): 组的数量，用于确定需要多少种颜色。
    """
    if palette_name == 'custom':
        # 自定义颜色列表
        custom_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        colors = custom_colors[:group_count]
    else:
        # 使用seaborn内置调色板
        colors = sns.color_palette(palette_name, group_count)
    
    # 创建组名到颜色的映射
    analysis.color_map = {group: color for group, color in zip(analysis.groups, colors)}
    # 更新palette属性以保持向后兼容
    analysis.palette = colors


# 示例数据生成函数（用于演示）
def generate_sample_data(analysis_type='time', include_cell_lines=False):
    """生成示例细胞活性数据。

    根据指定的分析类型生成模拟的细胞活性数据，可选择包含多个细胞株。
    - 时间依赖性数据：模拟不同时间点的细胞活性变化。
    - 剂量依赖性数据：使用四参数逻辑斯蒂方程模拟不同浓度下的细胞活性。

    Args:
        analysis_type (str, optional): 分析类型，'time'（时间依赖性）或'dose'（剂量依赖性）。
            默认为'time'。
        include_cell_lines (bool, optional): 是否包含多个细胞株的数据。如果为True，
            将生成3个不同敏感性的细胞株数据。默认为False。

    Returns:
        pd.DataFrame: 包含模拟细胞活性数据的DataFrame，包含以下列：
            - well_id: 样本ID（如'Cell_01_Compound_A_Rep1'）
            - compound: 化合物名称
            - time_hour/concentration_um: 时间或浓度值
            - viability_percentage: 细胞活性百分比
            - cell_line: 细胞株名称（如果include_cell_lines=True）
    """
    if analysis_type == 'time':
        # 生成时间依赖性数据
        samples = []
        # 定义时间点
        time_points = [0, 24, 48, 72, 96]
        # 定义化合物
        compounds = ['Control', 'Compound_A', 'Compound_B', 'Compound_C']
        # 为每个化合物定义活性下降模式
        decay_rates = {
            'Control': 0.001,       # 几乎不变
            'Compound_A': 0.02,     # 中等抑制
            'Compound_B': 0.05,     # 强抑制
            'Compound_C': 0.01      # 弱抑制
        }
        
        # 定义细胞株（如果包含）
        cell_lines = ['Cell_01', 'Cell_02', 'Cell_03'] if include_cell_lines else ['Cell_01']
        
        # 为每个化合物生成3个复孔的数据，如果包含细胞株则为每个细胞株生成数据
        for cell_line in cell_lines:
            for comp in compounds:
                # 为每个细胞株调整衰减率以模拟不同敏感性
                cell_factor = {'Cell_01': 1.0, 'Cell_02': 1.2, 'Cell_03': 0.8}.get(cell_line, 1.0)
                adjusted_decay = decay_rates[comp] * cell_factor
                
                for rep in range(1, 4):
                    well_id = f'{cell_line}_{comp}_Rep{rep}'
                    for time in time_points:
                        # 基础活性值（随时间衰减）
                        base_viability = 100 * np.exp(-adjusted_decay * time)
                        # 添加一些随机噪声
                        noise = np.random.normal(0, 3)
                        viability = max(0, min(100, base_viability + noise))
                        
                        row_data = {
                            'well_id': well_id,
                            'compound': comp,
                            'time_hour': time,
                            'viability_percentage': viability
                        }
                        
                        if include_cell_lines:
                            row_data['cell_line'] = cell_line
                        
                        samples.append(row_data)
        
        return pd.DataFrame(samples)
    
    else:  # dose
        # 定义剂量参数
        max_conc = 100
        drug_params = {
            'Drug_X': {'ic50': 1, 'hill': 1.2, 'min_viability': 5},
            'Drug_Y': {'ic50': 5, 'hill': 1.0, 'min_viability': 10},
            'Drug_Z': {'ic50': 20, 'hill': 0.8, 'min_viability': 20}
        }
        
        # 生成剂量依赖性数据
        samples = []
        # 定义浓度梯度（μM）
        concentrations = [0.001, 0.01, 0.1, 1, 10, 100]
        # 定义化合物
        compounds = ['Control', 'Drug_X', 'Drug_Y', 'Drug_Z']
        
        # 定义细胞株（如果包含）
        cell_lines = ['Cell_01', 'Cell_02', 'Cell_03'] if include_cell_lines else ['Cell_01']
        
        # 为每个细胞株定义敏感性因子
        cell_sensitivity = {
            'Cell_01': 1.0,  # 中等敏感性
            'Cell_02': 1.5,  # 更敏感
            'Cell_03': 0.6   # 更耐药
        }
        
        # 为每个化合物生成3个复孔的数据
        for cell_line in cell_lines:
            for drug in compounds:
                for rep in range(1, 4):
                    well_id = f'{cell_line}_{drug}_Rep{rep}'
                    for conc in concentrations:
                        if drug == 'Control':
                            # 对照组活性保持在100%左右
                            viability = max(0, min(100, 100 + np.random.normal(0, 5)))
                        else:
                            # 使用四参数逻辑斯蒂方程计算活性
                            params = drug_params[drug]
                            # 根据细胞株敏感性调整IC50
                            adjusted_ic50 = params['ic50'] / cell_sensitivity.get(cell_line, 1.0)
                            # 四参数方程
                            viability = params['min_viability'] + \
                                        (100 - params['min_viability']) / \
                                        (1 + (conc / adjusted_ic50)**params['hill'])
                            # 添加一些随机噪声
                            noise = np.random.normal(0, 4)
                            viability = max(0, min(100, viability + noise))
                        
                        row_data = {
                            'well_id': well_id,
                            'compound': drug,
                            'concentration_um': conc,
                            'viability_percentage': viability
                        }
                        
                        if include_cell_lines:
                            row_data['cell_line'] = cell_line
                        
                        samples.append(row_data)
        
        return pd.DataFrame(samples)


if __name__ == '__main__':
    # --- cellauc包的演示 ---
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='使用AUC方法进行细胞活性分析')
    parser.add_argument('--input', type=str, default=None, 
                        help='包含细胞活性数据的输入CSV文件路径。如果未提供，将生成示例数据。')
    parser.add_argument('--type', type=str, default='time', 
                        help='分析类型: time（时间依赖性）或dose（剂量依赖性）')
    parser.add_argument('--fig', type=str, default='box', 
                        help='绘图类型，可选值: box(箱型图), bar(柱状图), strip(散点图), violin(小提琴图), swarm(蜂群图)')
    parser.add_argument('--palette', type=str, default='husl', 
                        help='颜色方案，可选值: husl, tab10, Set2, viridis, plasma等')
    parser.add_argument('--by-cell-line', action='store_true', 
                        help='如果数据包含cell_line列，则按细胞株分组分析')
    parser.add_argument('--sample-data-with-cell-lines', action='store_true',
                        help='生成示例数据时包含多个细胞株')
    args = parser.parse_args()

    # 创建结果存储目录（使用时间戳）
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    result_dir = os.path.join('result', timestamp)
    os.makedirs(result_dir, exist_ok=True)
    print(f"结果将保存到: {result_dir}")

    # 1. 加载或生成数据
    if args.input:
        try:
            df = pd.read_csv(args.input)
            print(f"已加载数据文件: {args.input}")
        except FileNotFoundError:
            print(f"错误: '{args.input}' 文件未找到。生成示例数据...")
            df = generate_sample_data(args.type)
    else:
        print("未提供输入文件。生成示例数据...")
        df = generate_sample_data(
            args.type, 
            include_cell_lines=args.sample_data_with_cell_lines
        )
        # 保存示例数据用于参考
        sample_data_path = os.path.join(result_dir, f'sample_cell_data_{args.type}.csv')
        df.to_csv(sample_data_path, index=False)
        print(f"示例数据已保存到: {sample_data_path}")
        if args.sample_data_with_cell_lines:
            print("示例数据包含多个细胞株。使用--by-cell-line参数可按细胞株分析。")

    # 2. 初始化分析对象
    # 检查数据中是否包含cell_line列
    has_cell_line = any('cell' in col.lower() for col in df.columns)
    
    if args.type == 'time':
        # 时间依赖性分析
        analysis = CellViabilityAnalysis(
            data=df,
            sample_id_col='well_id',
            group_col='compound',
            x_col='time_hour',
            response_col='viability_percentage',
            analysis_type='time',
            cell_line_col='cell_line' if has_cell_line else None
        )
        # 设置分析范围（使用数据中的实际最小和最大值）
        min_time = df['time_hour'].min()
        max_time = df['time_hour'].max()
        total_range = (min_time, max_time)
    else:
        # 剂量依赖性分析
        analysis = CellViabilityAnalysis(
            data=df,
            sample_id_col='well_id',
            group_col='compound',
            x_col='concentration_um',
            response_col='viability_percentage',
            analysis_type='dose',
            cell_line_col='cell_line' if has_cell_line else None
        )
        # 设置分析范围
        min_conc = df['concentration_um'].min()
        max_conc = df['concentration_um'].max()
        total_range = (min_conc, max_conc)
    
    # 如果数据包含细胞株信息，显示可用的细胞株数量
    if analysis.cell_line_col:
        cell_line_count = df[analysis.cell_line_col].nunique()
        print(f"数据包含 {cell_line_count} 个细胞株")
        if args.by_cell_line:
            print("将按细胞株分组进行AUC分析")
    
    # 3. 设置颜色方案
    group_count = len(analysis.groups)
    set_palette(analysis, args.palette, group_count)

    # 4. 绘制原始响应曲线
    print("生成响应曲线图表...")
    curve_title = 'time_dependent_response_curves' if args.type == 'time' else 'dose_response_curves'
    curve_path = os.path.join(result_dir, f'{curve_title}.png')
    analysis.plot_response_curves(
        save_path=curve_path,
        log_scale=(args.type == 'dose')  # 剂量响应使用对数刻度
    )
    plt.close()  # 关闭图形以释放内存

    # 5. 计算总AUC
    print(f"计算总AUC（范围: {total_range}）...")
    total_auc_results = analysis.calculate_auc(total_x_range=total_range, by_cell_line=args.by_cell_line)
    print(f"AUC计算完成。共分析 {len(total_auc_results)} 个样本")
    
    # 如果按细胞株分析，显示前几行结果
    if args.by_cell_line and analysis.cell_line_col:
        print("\nAUC结果预览（前10行）:")
        print(total_auc_results.head(10))
    else:
        print(total_auc_results)

    # 6. 绘制总AUC汇总图并与'Control'组进行统计比较
    print("生成总AUC汇总图与统计分析...")
    
    # 根据是否按细胞株分析选择不同的绘图策略
    if args.by_cell_line and analysis.cell_line_col:
        # 如果按细胞株分析，生成热力图和分组图
        # 1. 热力图
        heatmap_path = os.path.join(result_dir, 'total_auc_heatmap.png')
        fig, ax = plt.subplots(figsize=(12, 16))  # 为大量细胞株设置较大的图形大小
        analysis.plot_auc_summary(
            metric='total_auc',
            plot_type='heatmap',
            by_cell_line=True,
            ax=ax,
            save_path=heatmap_path,
            cmap='viridis'
        )
        plt.tight_layout()
        plt.savefig(heatmap_path, dpi=300)
        plt.close()
        print(f"热力图已保存到: {heatmap_path}")
        
        # 2. 按细胞株分组的条形图（仅显示部分细胞株以避免过度拥挤）
        cell_lines = df[analysis.cell_line_col].unique()
        if len(cell_lines) > 20:
            # 如果细胞株太多，只显示前20个
            selected_cell_lines = list(cell_lines[:20])
            barplot_path = os.path.join(result_dir, 'total_auc_by_cell_line_top20.png')
            print(f"细胞株数量较多，仅显示前20个细胞株的条形图")
        else:
            selected_cell_lines = list(cell_lines)
            barplot_path = os.path.join(result_dir, f'total_auc_by_cell_line_{args.fig}.png')
        
        analysis.plot_auc_summary(
            metric='total_auc',
            plot_type='bar',
            by_cell_line=True,
            selected_cell_lines=selected_cell_lines,
            save_path=barplot_path
        )
        plt.close()
        print(f"按细胞株分组的条形图已保存到: {barplot_path}")
    else:
        # 标准汇总图
        total_auc_plot_path = os.path.join(result_dir, f'total_auc_comparison_{args.fig}.png')
        analysis.plot_auc_summary(
            metric='total_auc', 
            control_group=None,
            plot_type=args.fig,
            save_path=total_auc_plot_path
        )
        plt.close()  # 关闭图形以释放内存
        print(f"汇总图已保存到: {total_auc_plot_path}")
    
    # 7. 如果按细胞株分析，额外生成按化合物分组的汇总图
    if args.by_cell_line and analysis.cell_line_col:
        compound_summary_path = os.path.join(result_dir, f'total_auc_by_compound_{args.fig}.png')
        analysis.plot_auc_summary(
            metric='total_auc',
            control_group=None,
            plot_type=args.fig,
            save_path=compound_summary_path
        )
        plt.close()
        print(f"按化合物分组的汇总图已保存到: {compound_summary_path}")

    
    # 10. 生成并保存样本个体AUC数据CSV文件
    print("\n生成个体样本AUC数据CSV...")
    merged_auc_data = total_auc_results
    # 保存到CSV
    individual_auc_csv_path = os.path.join(result_dir, 'individual_sample_auc_data.csv')
    merged_auc_data.to_csv(individual_auc_csv_path, index=False)
    print(f"个体样本AUC数据已保存到: {individual_auc_csv_path}")
    
    # 11. 计算并保存小组平均AUC数据CSV文件
    print("\n生成小组平均AUC数据CSV...")
    # 获取总AUC列
    auc_columns = ['total_auc']
    
    # 根据是否按细胞株分析选择不同的分组方式
    if args.by_cell_line and analysis.cell_line_col:
        # 按化合物分组计算平均AUC
        compound_avg_auc = merged_auc_data.groupby(analysis.group_col)[auc_columns].mean().reset_index()
        compound_avg_auc.columns = [f"{col}_mean" if col in auc_columns else col for col in compound_avg_auc.columns]
        compound_avg_csv_path = os.path.join(result_dir, 'compound_average_auc_data.csv')
        compound_avg_auc.to_csv(compound_avg_csv_path, index=False)
        print(f"化合物平均AUC数据已保存到: {compound_avg_csv_path}")
        
        # 按细胞株分组计算平均AUC
        cell_line_avg_auc = merged_auc_data.groupby(analysis.cell_line_col)[auc_columns].mean().reset_index()
        cell_line_avg_auc.columns = [f"{col}_mean" if col in auc_columns else col for col in cell_line_avg_auc.columns]
        cell_line_avg_csv_path = os.path.join(result_dir, 'cell_line_average_auc_data.csv')
        cell_line_avg_auc.to_csv(cell_line_avg_csv_path, index=False)
        print(f"细胞株平均AUC数据已保存到: {cell_line_avg_csv_path}")
        
        # 保存完整的细胞株×化合物矩阵
        pivot_auc = merged_auc_data.pivot(index=analysis.cell_line_col, 
                                         columns=analysis.group_col, 
                                         values='total_auc')
        matrix_csv_path = os.path.join(result_dir, 'cell_line_compound_auc_matrix.csv')
        pivot_auc.to_csv(matrix_csv_path)
        print(f"细胞株×化合物AUC矩阵已保存到: {matrix_csv_path}")
    else:
        # 原始行为：按组计算平均值
        group_avg_auc = merged_auc_data.groupby(analysis.group_col)[auc_columns].mean().reset_index()
        group_avg_auc.columns = [f"{col}_mean" if col in auc_columns else col for col in group_avg_auc.columns]
        group_avg_auc_csv_path = os.path.join(result_dir, 'group_average_auc_data.csv')
        group_avg_auc.to_csv(group_avg_auc_csv_path, index=False)
        print(f"小组平均AUC数据已保存到: {group_avg_auc_csv_path}")
    
    print("\n所有分析和图表生成完成！")